from align import Table,SubstMatrix,smith_Waterman,create_submat

def align_query(query, seqs, sm:dict, gap):
    def socre(table):
        cell = table.smith_Waterman(sm,gap)[-1]
        return cell.value
    sm = SubstMatrix().build_from_dict(sm)
    table = max((Table(query,seq) for seq in seqs),key=socre)
    return table.recover_align_local()[-1]

def read_database(filename):
    with open(filename, 'r') as f:
        db = f.readlines()
    return [line.rstrip() for line in db]

def build_map(query, word_size):
    word_map = dict()
    for i in range(len(query)-word_size+1):
        subseq = query[i:i+word_size]
        if subseq in word_map:
            word_map[subseq].append(i)
        else:
            word_map[subseq] = [i]
    return word_map

def get_hits(query, word_map, word_size):
    hits = []
    for i in range(len(query)-word_size+1):
        subseq = query[i:i+word_size]
        if subseq in word_map:
            tag_indexs = word_map[subseq]
            for index in tag_indexs:
                hits.append((index,i))
    return hits

def extend_hit(seq, hit:tuple, query, word_size) -> tuple[int]:
    '''
    argument: str, (int,int), str, int\n
    return tuple:\n
    1. seq match start index\n
    2. query match start index\n
    3. match length\n
    4. match base num 
    '''
    stq, sts = hit[0], hit[1]
    ## move forward
    match_fw = 0       
    k=0
    bestk = 0
    while (2*match_fw >= k and 
           stq+word_size+k < len(query) and 
           sts+word_size+k < len(seq)):
        if query[stq+word_size+k] == seq[sts+word_size+k]: 
            match_fw+=1
            bestk = k+1
        k += 1
    size = word_size + bestk
    ## move backwards
    k = 0
    match_bw = 0   
    bestk = 0
    while 2*match_bw >= k and stq > k and sts > k:
        if query[stq-k-1] == seq[sts-k-1]: 
            match_bw+=1
            bestk = k+1
        k+=1
    size += bestk
    
    return (stq-bestk, sts-bestk, size, word_size+match_fw+match_bw) 

def hit_best_score(seq, query, word_map, word_size) -> tuple[int]:
    hits = get_hits(seq, word_map, word_size)
    bestScore = -1.0
    best = ()
    for h in hits:
        ext = extend_hit(seq, h, query, word_size)
        score = ext[3]
        if score > bestScore or (score == bestScore and ext[2] < best[2]):
            bestScore = score
            best = ext
    return best

def bset_alignment(db,query,word_size):
    '''
    return tuple:\n
    1-4: return of extend_hit\n
    5. k-th seq in database\n
    '''
    word_map = build_map(query, word_size)
    bestScore = -1.0
    res = tuple()
    for k in range(len(db)):    # k-th seq in database
        bestSeq = hit_best_score(db[k], query, word_map, word_size)
        if bestSeq != ():
            score = bestSeq[3]  
            if score > bestScore or (score == bestScore and bestSeq[2] < res[2]):
                bestScore = score
                res = (bestSeq[0], bestSeq[1], bestSeq[2], bestSeq[3], k)
    if bestScore < 0:
        return tuple()
    else:
        return res

if __name__ == "__main__":
    sm = create_submat(1,-1,"ATCG")
    print(align_query('ATT',['ATTC','CGCG','TAG'],sm,-1))
    print(read_database('bioinformatic_algorithm_colg/5.Search_Databases/test_db'))
    print(build_map('ATCTCATC',2))
    pass